from fastapi import FastAPI, Body, File, UploadFile
from paddleocr import PaddleOCR, draw_ocr
import re
import io  
from PIL import Image

app = FastAPI()


def get_text(url: str) -> []:
    ocr = PaddleOCR(use_angle_cls=True, lang="ch")
    result = ocr.ocr(url, cls=True)
    data = []
    for idx in range(len(result)):
        res = result[idx]
        for line in res:
            # print(line[1])
            data.append({"content": line[1][0], "rate": line[1][1]})
    return data


@app.get("/")
async def root():
    return {"message": "Welcome to OCR system.2"}


@app.post("/ocr")
async def ocr(url: str = Body(..., title='image url', embed=True),
              content: str = Body(..., title='ocr text content', embed=True)):
    data = get_text(url)
    match = 0
    rate = ""
    for idx in range(len(data)):
        ocrContent = data[idx]["content"]
        ocrContent = ocrContent.strip()
        ocrContent = str.replace(ocrContent, ',', '')
        print(ocrContent)
        if ocrContent.find(content) != -1:
            rate = data[idx]["rate"]
            match = 1
            break
    return {"code": 200, "match": match, "rate": rate, "ocr": data}


@app.post("/ocr2")
async def ocr2(file: UploadFile = File(...)):
    image_bytes = await file.read()  
    image_stream = io.BytesIO(image_bytes)  
      
    # 使用 PIL 打开图片  
    image = Image.open(image_stream).convert('RGB')  
      
    # 使用 PaddleOCR 进行 OCR  
    result = ocr.ocr(image, use_gpu=False)  # 设置 use_gpu 为 True 如果你的机器支持 GPU  
      
    data = []
    for idx in range(len(result)):
        res = result[idx]
        for line in res:
            # print(line[1])
            data.append({"content": line[1][0], "rate": line[1][1]})
    return {"code": 200,  "ocr": data}